#!/usr/bin/env python
#
# Copyright 2012,2013,2015,2020 Free Software Foundation, Inc.
#
# This file is part of GNU Radio
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
#

"""
Read samples from a UHD device and write to file formatted as binary
outputs single precision complex float values or complex short values
(interleaved 16 bit signed short integers).
"""

import sys
import os
import pmt
import time
from gnuradio import gr, eng_notation
from gnuradio import blocks
from gnuradio import uhd
from gnuradio.eng_option import eng_option
from optparse import OptionParser

n2s = eng_notation.num_to_str
COMMAND_DELAY = .2


class rx_cfile_block(gr.top_block):
    """
    Simple flowgraph that streams from USRP source to file.
    """

    def __init__(self, options, filename):
        gr.top_block.__init__(self)
        # Dissect the channel list:
        try:
            self.channels = [int(x.strip())
                             for x in options.channels.split(",")]
        except ValueError:
            sys.stderr.write(
                "[UHD_RX] [ERROR] Invalid channel list: {}".format(options.channels))
            exit(1)
        # Create file names:
        if len(self.channels) == 1:
            self.filenames = [filename, ]
        else:
            base, ext = os.path.splitext(filename)
            self.filenames = []
            for i in range(len(self.channels)):
                self.filenames.append(
                    "{base}.{num}{ext}".format(base=base, num=i, ext=ext))
        # Check CPU format:
        self.cpu_format = 'fc32'
        self.item_size = gr.sizeof_gr_complex
        self.meta_file_type = blocks.GR_FILE_FLOAT
        if options.output_shorts:
            self.cpu_format = 'sc16'
            self.item_size = gr.sizeof_short * 2
            self.meta_file_type = blocks.GR_FILE_SHORT
        # Create a UHD device source:
        self._u = uhd.usrp_source(
            device_addr=options.args,
            stream_args=uhd.stream_args(
                self.cpu_format,
                options.wire_format,
                args=options.stream_args,
                channels=self.channels,
            )
        )
        # Set the subdevice spec:
        if options.spec:
            for mb_idx in range(self._u.get_num_mboards()):
                self._u.set_subdev_spec(options.spec, mb_idx)
        # Set the antenna:
        if options.antenna is not None:
            self.antenna = [x.strip() for x in options.antenna.split(",")]
            if len(self.antenna) != 1 and len(self.antenna) != len(self.channels):
                sys.stderr.write("[UHD_RX] [ERROR] Invalid antenna setting for {} channels: {}".format(
                    len(self.channels), options.antenna
                ))
                exit(1)
            if len(self.antenna) == 1 and len(self.channels) > 1:
                self.antenna = [self.antenna[0], ] * len(self.channels)
            for i, chan in enumerate(self.channels):
                self._u.set_antenna(self.antenna[i], chan)
                if options.verbose:
                    print("[UHD_RX] Channel {chan}: Using antenna {ant}.".format(
                        chan=chan, ant=self._u.get_antenna(chan)
                    ))
        # Set receiver sample rate:
        self._u.set_samp_rate(options.samp_rate)
        samp_rate = self._u.get_samp_rate()
        # Set receive daughterboard gain:
        if options.gain is None:
            print("[UHD_RX] Defaulting to mid-point gains:")
            for chan in self.channels:
                self._u.set_normalized_gain(.5, chan)
                print("[UHD_RX] Channel {chan} gain: {g} dB".format(
                    chan=chan, g=self._u.get_gain(chan)))
        else:
            for chan in self.channels:
                if options.normalized_gain:
                    self._u.set_normalized_gain(options.gain, chan)
                else:
                    self._u.set_gain(options.gain, chan)
        gain = self._u.get_gain(self.channels[0])
        # Set frequency (tune request takes lo_offset):
        if options.lo_offset is not None:
            treq = uhd.tune_request(options.freq, options.lo_offset)
        else:
            treq = uhd.tune_request(options.freq)
        # Make sure tuning is synched:
        command_time_set = False
        if len(self.channels) > 1:
            if options.sync == 'pps':
                self._u.set_time_unknown_pps(uhd.time_spec())
            cmd_time = self._u.get_time_now() + uhd.time_spec(COMMAND_DELAY)
            try:
                for mb_idx in range(self._u.get_num_mboards()):
                    self._u.set_command_time(cmd_time, mb_idx)
                command_time_set = True
            except RuntimeError:
                sys.stderr.write(
                    '[UHD_RX] [WARNING] Failed to set command times.\n')
        for chan in self.channels:
            tr = self._u.set_center_freq(treq, chan)
            if tr == None:
                sys.stderr.write(
                    '[UHD_RX] [ERROR] Failed to set center frequency on channel {chan}\n'.format(chan=chan))
                exit(1)
        if command_time_set:
            for mb_idx in range(self._u.get_num_mboards()):
                self._u.clear_command_time(mb_idx)
            print("[UHD_RX] Syncing channels...")
            time.sleep(COMMAND_DELAY)
        freq = self._u.get_center_freq(self.channels[0])
        # Create file sink(s):
        self._sink = []
        for i in range(len(self.channels)):
            if options.metafile:
                # store additional metadata
                extras = pmt.make_dict()
                extras = pmt.dict_add(extras, pmt.intern(
                    "rx_gain"), pmt.from_double(gain))
                extras_str = pmt.serialize_str(extras)
                self._sink.append(blocks.file_meta_sink(
                    self.item_size, self.filenames[i],
                    samp_rate, 1,
                    self.meta_file_type, True,
                    1000000, extras_str, False
                ))
            else:
                self._sink.append(blocks.file_sink(
                    self.item_size, self.filenames[i]))
            # Create head block if needed and wire it up:
            if options.nsamples is None:
                self.connect((self._u, i), self._sink[i])
            else:
                self._head = blocks.head(self.item_size, int(options.nsamples))
                self.connect((self._u, i), self._head, self._sink[i])
        # Output status info if requested:
        if options.verbose:
            try:
                info = self._u.get_usrp_info()
                mboard_id = info["mboard_id"].split(" ")[0]
                if info["mboard_serial"] == "":
                    mboard_serial = "no serial"
                else:
                    mboard_serial = info["mboard_serial"]
                rx_id = info["rx_id"].split(" ")[0]
                if info["rx_serial"] == "":
                    rx_serial = "no serial"
                else:
                    rx_serial = info["rx_serial"]
                rx_antenna = info["rx_antenna"]
                rx_subdev_spec = info["rx_subdev_spec"]
                print("[UHD_RX] Motherboard: %s (%s)" %
                      (mboard_id, mboard_serial))
                if "B200" in mboard_id or "B210" in mboard_id or "E310" in mboard_id:
                    print("[UHD_RX] Daughterboard: %s (%s, %s)" %
                          (mboard_id, rx_antenna, rx_subdev_spec))
                else:
                    print("[UHD_RX] Daughterboard: %s (%s, %s, %s)" %
                          (rx_id, rx_serial, rx_antenna, rx_subdev_spec))
            except KeyError:
                print("[UHD_RX] Args: ", options.args)
            print("[UHD_RX] Receiving on {} channels.".format(len(self.channels)))
            print("[UHD_RX] Rx gain:               {gain}".format(gain=gain))
            print("[UHD_RX] Rx frequency:          {freq}".format(freq=freq))
            print("[UHD_RX] Rx baseband frequency: {actual}".format(
                actual=n2s(tr.actual_rf_freq)))
            print("[UHD_RX] Rx DDC frequency:      {dsp}".format(
                dsp=n2s(tr.actual_dsp_freq)))
            print("[UHD_RX] Rx Sample Rate:        {rate}".format(
                rate=n2s(samp_rate)))
            if options.nsamples is None:
                print("[UHD_RX] Receiving samples until Ctrl-C")
            else:
                print("[UHD_RX] Receiving {n} samples.".format(
                    n=n2s(options.nsamples)))
            if options.output_shorts:
                print("[UHD_RX] Writing 16-bit complex shorts")
            else:
                print("[UHD_RX] Writing 32-bit complex floats")
            print("[UHD_RX] Output file(s): {files}".format(
                files=", ".join(self.filenames)))
        # Direct asynchronous notifications to callback function:
        if options.show_async_msg:
            self.async_msgq = gr.msg_queue(0)
            self.async_src = uhd.amsg_source("", self.async_msgq)
            self.async_rcv = uhd.msgq_runner(
                self.async_msgq, self.async_callback)

    def async_callback(self, msg):
        md = self.async_src.msg_to_async_metadata_t(msg)
        print("[UHD_RX] Channel: %i Time: %f Event: %i" %
              (md.channel, md.time_spec.get_real_secs(), md.event_code))


def get_options():
    usage = "%prog: [options] output_filename"
    parser = OptionParser(option_class=eng_option, usage=usage)
    parser.add_option("-a", "--args", type="string", default="",
                      help="UHD device address args , [default=%default]")
    parser.add_option("", "--spec", type="string", default=None,
                      help="Subdevice of UHD device where appropriate")
    parser.add_option("-c", "--channels", type="string", default="0",
                      help="Select receive channels")
    parser.add_option("-A", "--antenna", type="string", default=None,
                      help="Select Rx Antenna(s) where appropriate.\nUse a comma-delimited list if different channels have different antenna ports.")
    parser.add_option("-r", "--samp-rate", type="eng_float", default=1e6,
                      help="Set sample rate (bandwidth) [default=%default]")
    parser.add_option("-f", "--freq", type="eng_float", default=None,
                      help="Set frequency to FREQ", metavar="FREQ")
    parser.add_option("", "--lo-offset", type="eng_float", default=None,
                      help="Set daughterboard LO offset to OFFSET [default=hw default]")
    parser.add_option("-g", "--gain", type="eng_float", default=None,
                      help="Set gain in dB (default is midpoint)")
    parser.add_option("--normalized-gain", action="store_true",
                      help="Specify gain as normalized value (in [0, 1])")
    parser.add_option("-m", "--metafile", action="store_true", default=False,
                      help="output metadata to file [default=%default]")
    parser.add_option("-s", "--output-shorts", action="store_true", default=False,
                      help="Output interleaved shorts instead of complex floats")
    parser.add_option("-N", "--nsamples", type="eng_float", default=None,
                      help="Number of samples to collect [default=+inf]")
    parser.add_option("-v", "--verbose", action="store_true", default=False,
                      help="verbose output")
    parser.add_option("", "--wire-format", type="string", default="sc16",
                      help="Set wire format from USRP [default=%default")
    parser.add_option("", "--stream-args", type="string", default="",
                      help="Set additional stream arguments")
    parser.add_option("", "--show-async-msg", action="store_true", default=False,
                      help="Show asynchronous message notifications from UHD [default=%default]")
    parser.add_option("", "--sync", type="choice", choices=('default', 'pps'),
                      default='default', help="Set to 'pps' to sync devices to PPS instead of internal.")
    (options, args) = parser.parse_args()
    if len(args) != 1:
        parser.print_help()
        exit(1)
    if options.freq is None:
        parser.print_help()
        sys.stderr.write('You must specify the frequency with -f FREQ\n')
        exit(1)
    return (options, args[0])


if __name__ == '__main__':
    (options, filename) = get_options()
    tb = rx_cfile_block(options, filename)
    try:
        tb.run()
    except KeyboardInterrupt:
        print("Receive terminated.")
